import os
import cv2 as cv
from matplotlib import pyplot as plt


# 工具方法

# red #ff0000
# green #00ff00
# 处理画好的k线图，便于后续CNN使用
def imageProcess(filename):
    src = cv.imread(filename)  # 读取图片
    height, width, channel = src.shape
    for row in range(height):
        for col in range(width):
            (b, g, r) = src[row, col]  # 获取像素点的RGB值
            if 200 <= r <= 255 and g < 50 and b < 50:  # red
                src[row, col] = (0, 0, 255)  # b,g,r
            elif r < 50 and 200 <= g <= 255 and b < 50:  # green
                src[row, col] = (0, 255, 0)  # b,g,r
            else:
                src[row, col] = (255, 255, 255)
            pass
    # src = cv.cvtColor(src, cv.COLOR_BGR2GRAY)  # 转为灰度
    cv.imwrite(filename, src)  # 保存图片


# 读取文件中的元素列表
def read_file_list(foldername, filename):
    with open(foldername + filename, 'r') as f:
        element_list = f.readlines()
    for i in range(len(element_list)):
        element_list[i] = element_list[i].rstrip()  # 舍掉右边的\n
    return element_list


# 创建文件夹
def make_folder(folder_path):
    '''make a folder if the folder not exist'''
    if os.path.exists(folder_path) is False:  # 若判断路径不存在
        os.mkdir(folder_path)  # 创建文件夹


# 股票代码标准化
def stock_code_standardization(stock):
    if stock[0] == '6':  # 沪市
        return stock + '.SH'
    else:  # 深市
        return stock + '.SZ'


# 绘制loss曲线
def plot_curve(data_x, data_list, title, xlabel, ylabel):
    fig = plt.figure()
    plt.plot(data_x, data_list, color='blue', marker="o", markersize=6)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(title + ".jpg")  # 保存图片 bbox_inches = 'tight'
